# -*- coding: utf-8 -*-

from transformers import BertTokenizer

class ChineseBERTTokenizerWrapper:
    def __init__(self, model_name="bert-base-chinese"):
        """初始化中文 BERT 分词器"""
        self.tokenizer = BertTokenizer.from_pretrained(model_name)

    def tokenize(self, text):
        """分词函数，将中文文本分解为 BERT 的子词片段"""
        tokens = self.tokenizer.tokenize(text)
        return tokens

    def encode(self, text, add_special_tokens=True):
        """编码函数，将中文文本转换为词汇表的索引"""
        token_ids = self.tokenizer.encode(text, add_special_tokens=add_special_tokens)
        return token_ids

    def decode(self, token_ids):
        """解码函数，将词汇表索引转换回文本"""
        text = self.tokenizer.decode(token_ids)
        return text

    def get_special_tokens(self):
        """获取中文 BERT 模型的特殊标记 (CLS 和 SEP)"""
        cls_token = self.tokenizer.cls_token
        sep_token = self.tokenizer.sep_token
        return cls_token, sep_token

# 使用示例
if __name__ == "__main__":
    text = "我爱自然语言处理和机器学习。"

    # 创建中文 BERT 分词器封装类实例
    tokenizer = ChineseBERTTokenizerWrapper()

    # 1. 分词
    tokens = tokenizer.tokenize(text)
    print("分词结果：", tokens)

    # 2. 编码
    token_ids = tokenizer.encode(text)
    print("编码结果：", token_ids)

    # 3. 解码
    decoded_text = tokenizer.decode(token_ids)
    print("解码结果：", decoded_text)

    # 4. 获取特殊标记
    cls_token, sep_token = tokenizer.get_special_tokens()
    print("CLS 标记：", cls_token)
    print("SEP 标记：", sep_token)

